# 导入所需的类和模块
from langchain.agents import create_react_agent, AgentExecutor
from langchain_community.llms.tongyi import Tongyi
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_core.tools import Tool

# 初始化 TongYi 模型
llm = Tongyi()


def tool1(input):
    return


tools = [Tool(func=tool1, name="tool1", description="当我问你我是谁的时候，你可以告诉我，我叫什么名字")]
template = '''Answer the following questions as best you can. You have access to the following tools:

            {tools}

            Use the following format:

            Question: the input question you must answer
            Thought: you should always think about what to do
            Action: the action to take, should be one of [{tool_names}]
            Action Input: the input to the action
            Observation: the result of the action
            ... (this Thought/Action/Action Input/Observation can repeat N times)
            Thought: I now know the final answer
            Final Answer: the final answer to the original input question

            Begin!

            Question: {input}
            Thought:{agent_scratchpad}'''

prompt = ChatPromptTemplate.from_template(template)

agent = create_react_agent(prompt=prompt, tools=tools, llm=llm)

agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# 初始化对话历史
history = ChatMessageHistory()

runnable = RunnableWithMessageHistory(
    agent_executor,
    lambda session_id: history,
    input_messages_key="input",
    history_messages_key="chat_history",
)

while True:
    user_input = input("User: ")
    res = runnable.invoke({
        "input": user_input},
        config={"configurable": {"session_id": "session_id"}}
    )
    print(f"Agent: {res['output']}")
